import torch
from PIL import Image
import matplotlib.pyplot as plt
from torchvision import transforms
plt.rcParams['font.sans-serif'] = ['SimHei']  # Windows 系统黑体
plt.rcParams['axes.unicode_minus'] = False  # 解决负号显示问题

# 加载类别名称
with open(r"F:\人工智能教材编写\traffic_sign\classes.txt", "r", encoding="utf-8") as f:
    class_names = [line.strip() for line in f.readlines()]

# 定义预处理（与验证集相同）
transform = transforms.Compose([
    transforms.Resize(256),
    transforms.CenterCrop(224),
    transforms.ToTensor(),
    transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
])

# # 加载训练好的模型
model = torch.load('resnet18_traffic.pth')
model.eval()

# 读取测试图像
img_path = r"F:\人工智能教材编写\traffic_sign\test\05_00.png"
pil_img = Image.open(img_path).convert('RGB')

# 预处理并预测
input_tensor = transform(pil_img).unsqueeze(0)
with torch.no_grad():
    output = model(input_tensor)
pred_class = output.argmax(dim=1).item()
print(f"预测类别: {class_names[pred_class]}")

# 显示结果
# 修改显示部分
plt.figure(figsize=(8, 4))
plt.imshow(pil_img)
plt.title(f'预测结果: {class_names[pred_class]}', fontsize=12)
plt.axis('off')
# plt.tight_layout()
plt.show()



